package publish_confirm;

import com.rabbitmq.client.Channel;
import com.rabbitmq.client.ConfirmListener;
import com.rabbitmq.client.Connection;
import com.rabbitmq.client.ConnectionFactory;
import constant.Constants;

import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.SortedSet;
import java.util.TreeSet;
import java.util.concurrent.ConcurrentNavigableMap;
import java.util.concurrent.ConcurrentSkipListMap;
import java.util.concurrent.TimeoutException;

/**
 * 发布确认模式
 */
public class PublishConfirmDemo {

    /**
     * 发送的消息数量
     */
    private static final Integer MESSAGE_COUNT = 1000;

    private static Connection createConnection() throws IOException, TimeoutException {
        //建立连接
        ConnectionFactory connectionFactory = new ConnectionFactory();
        connectionFactory.setHost(Constants.HOST);
        connectionFactory.setPort(Constants.PORT);
        connectionFactory.setUsername(Constants.USERNAME);
        connectionFactory.setPassword(Constants.PASSWORD);
        connectionFactory.setVirtualHost(Constants.VIRTUAL_HOST);
        return connectionFactory.newConnection();
    }

    public static void main(String[] args) throws IOException, TimeoutException {

        //单独确认
        publishMessagesIndividually();

        //批量确认
        publishMessagesInBatch();

        //异步确认
        handlePublishConfirmsAsynchronously();
    }

    /**
     * 异步确认
     */
    private static void handlePublishConfirmsAsynchronously() {

        try (Connection connection = createConnection()) {
            //开启信道
            Channel channel = connection.createChannel();

            //声明交换机

            //设置信道为 confirm 模式
            channel.confirmSelect();

            //声明队列
            channel.queueDeclare(Constants.PUBLISH_CONFIRM_QUEUE3, true, false, false, null);

            long startTime = System.currentTimeMillis();
            //存放 deliveryTag 和 message 的集合
            //key 为 deliveryTag，value 为对应的 message
            /**
             * 选择 ConcurrentNavigableMap 而不是 ConcurrentHashMap：
             * 基于跳表实现，清除所有 <= deliveryTag 的时间复杂度为 o(logn)
             * 线程安全
             * 保证了原子性
             */
            ConcurrentNavigableMap<Long, String> map = new ConcurrentSkipListMap<>();

            channel.addConfirmListener(new ConfirmListener() {

                /**
                 * 消息成功确认
                 * @param deliveryTag
                 * @param multiple true 表示批量删除；false 表示值删除当前 deliveryTag
                 */
                @Override
                public void handleAck(long deliveryTag, boolean multiple) throws IOException {
                    //将 deliveryTag 及以下的从 sortedSet 中去除
                    if (multiple) {
                        //批量删除
                        //true 表示包含 deliveryTag
                        map.headMap(deliveryTag, true).clear();
                    } else {
                        //单独删除
                        map.remove(deliveryTag);
                    }
                }

                /**
                 * 消息确认失败
                 * @param deliveryTag
                 * @param multiple
                 * @throws IOException
                 */
                @Override
                public void handleNack(long deliveryTag, boolean multiple) throws IOException {

                    String message = (String) map.get(deliveryTag);

                    if (multiple) {
                        //批量删除
                        map.headMap(
                                deliveryTag, true //包含 deliveryTag
                        ).clear();
                    } else {
                        //单独删除
                        map.remove(deliveryTag);
                    }

                    //消息确认失败，重新发送该消息
                    deliveryTag = channel.getNextPublishSeqNo(); //每次生成的 deliveryTag 都不一样
                    channel.basicPublish("", Constants.PUBLISH_CONFIRM_QUEUE3, null, message.getBytes(StandardCharsets.UTF_8));
                    map.put(deliveryTag, message);
                }
            });

            //发送消息
            int maxUnConfirmed = 100; //未处理消息的阈值
            for (int i = 0; i < MESSAGE_COUNT; i++) {

                //将未处理的消息数量降到阈值以下
                while (map.size() >= maxUnConfirmed) {
                    Thread.sleep(10);
                }

                long deliveryTag = channel.getNextPublishSeqNo(); //每次生成的 deliveryTag 都不一样
                String message = "Asynchronous Confirmation " + i;
                channel.basicPublish("", Constants.PUBLISH_CONFIRM_QUEUE3, null, message.getBytes(StandardCharsets.UTF_8));
                map.put(deliveryTag, message);
            }

            //保证所有消息都已经确认
            while (!map.isEmpty()) {
                Thread.sleep(10);
            }

            long endTime = System.currentTimeMillis();
            System.out.println("异步确认模式耗时: " + (endTime - startTime) + "ms");

        } catch (IOException | TimeoutException | InterruptedException e) {
            throw new RuntimeException(e);
        }

    }


    /**
     * 异步确认
     */
    private static void handlePublishConfirmsAsynchronously1() {

        try (Connection connection = createConnection()) {
            //开启信道
            Channel channel = connection.createChannel();

            //声明交换机

            //设置信道为 confirm 模式
            channel.confirmSelect();

            //声明队列
            channel.queueDeclare(Constants.PUBLISH_CONFIRM_QUEUE3, true, false, false, null);

            long startTime = System.currentTimeMillis();
            //存放 deliveryTag 的集合
            SortedSet<Long> sortedSet = new TreeSet<>();

            channel.addConfirmListener(new ConfirmListener() {

                /**
                 * 消息成功确认
                 * @param deliveryTag
                 * @param multiple true 表示批量删除；false 表示值删除当前 deliveryTag
                 */
                @Override
                public void handleAck(long deliveryTag, boolean multiple) throws IOException {
                    //将 deliveryTag 及以下的从 sortedSet 中去除
                    if (multiple) {
                        //批量删除
                        sortedSet.headSet(deliveryTag + 1).clear();
                    } else {
                        //单独删除
                        sortedSet.remove(deliveryTag);
                    }
                }

                /**
                 * 消息确认失败
                 * @param deliveryTag
                 * @param multiple
                 * @throws IOException
                 */
                @Override
                public void handleNack(long deliveryTag, boolean multiple) throws IOException {
                    //将 deliveryTag 及以下的从 sortedSet 中去除
                    if (multiple) {
                        //批量删除
                        sortedSet.headSet(deliveryTag + 1).clear();
                    } else {
                        //单独删除
                        sortedSet.remove(deliveryTag);
                    }

                    //消息确认失败，根据具体业务采取措施
                }
            });

            //发送消息
            int maxUnConfirmed = 100; //未处理消息的阈值
            for (int i = 0; i < MESSAGE_COUNT; i++) {

                //将未处理的消息数量降到阈值以下
                while (sortedSet.size() >= maxUnConfirmed) {
                    Thread.sleep(10);
                }

                long deliveryTag = channel.getNextPublishSeqNo(); //每次生成的 deliveryTag 都不一样
                String message = "Asynchronous Confirmation " + i;
                channel.basicPublish("", Constants.PUBLISH_CONFIRM_QUEUE3, null, message.getBytes(StandardCharsets.UTF_8));
                sortedSet.add(deliveryTag);
            }

            //保证所有消息都已经确认
            while (!sortedSet.isEmpty()) {
                Thread.sleep(10);
            }

            long endTime = System.currentTimeMillis();
            System.out.println("异步确认模式耗时: " + (endTime - startTime) + "ms");

        } catch (IOException | TimeoutException | InterruptedException e) {
            throw new RuntimeException(e);
        }

    }

    /**
     * 批量确认
     */
    private static void publishMessagesInBatch() {
        try (Connection connection = createConnection()) {

            //开启信道
            Channel channel = connection.createChannel();

            //将信道设置为 confirm 模式
            channel.confirmSelect();

            //使用默认的交换机, routingKey 为队列名

            //声明队列
            channel.queueDeclare(Constants.PUBLISH_CONFIRM_QUEUE2, true, false, false, null);

            int batchMessageCount = 100;
            int sendMessageCount = 0;

            //发布消息
            long startTime = System.currentTimeMillis();
            for (int i = 0; i < MESSAGE_COUNT; i++) {
                String message = "Batch Confirmation " + i;
                channel.basicPublish("", Constants.PUBLISH_CONFIRM_QUEUE2, null, message.getBytes(StandardCharsets.UTF_8));

                sendMessageCount++;
                if (sendMessageCount == batchMessageCount) {
                    channel.waitForConfirmsOrDie(5000);
                    sendMessageCount = 0;
                }
            }

            //还有没有确认的消息
            if (sendMessageCount > 0) {
                channel.waitForConfirmsOrDie(5000);
            }

            long endTime = System.currentTimeMillis();
            System.out.println("批量确认模式耗时: " + (endTime - startTime) + "ms");

        } catch (IOException | TimeoutException | InterruptedException e) {
            throw new RuntimeException(e);
        }
    }

    /**
     * 单独确认
     */
    private static void publishMessagesIndividually() {
         try (Connection connection = createConnection()) {

             //开启信道
             Channel channel = connection.createChannel();

             //将信道设置为 confirm 模式
             channel.confirmSelect();

             //使用默认的交换机, routingKey 为队列名

             //声明队列
             channel.queueDeclare(Constants.PUBLISH_CONFIRM_QUEUE1, true, false, false, null);

             //发布消息
             long startTime = System.currentTimeMillis();
             for (int i = 0; i < MESSAGE_COUNT; i++) {
                 String message = "Individual confirmation " + i;
                 channel.basicPublish("", Constants.PUBLISH_CONFIRM_QUEUE1, null, message.getBytes(StandardCharsets.UTF_8));

                 //等待 broker 确认
                 //最长等待时间 5000ms
                 channel.waitForConfirmsOrDie(5000);
             }
             long endTime = System.currentTimeMillis();
             System.out.println("单独确认模式耗时: " + (endTime - startTime) + "ms");

         } catch (IOException | TimeoutException | InterruptedException e) {
             throw new RuntimeException(e);
         }
    }
}
